SageMakerのスペースとエンドポイントの停止忘れをチェックするLambdaを作ってみた

#Amazon SageMaker

#SageMaker Studio

よなみね

2024.03.25

はじめに

みなさんこんにちは、先日SageMakerのエンドポイントを消し忘れて200ドルの請求がきたよなです。
今回はエンドポイントと合わせて、Studio内で起動したスペースの停止忘れをチェックするためのLambdaを作成してみました。

エンドポイントのチェックは他の記事を参考にすぐに実装できたのでコードのみ記載しています。
今回は特にスペースのステータスを確認することにフォーカスした記事を書いています。

構成図

今回作成する構成図はこのような感じです。

全体の流れはこんな感じです。

EventBridgeで定期的にLambdaを実行
Lambdaでリソースのステータスをチェック
SNSでステータスをメール通知

SageMakerのリソースは学習を実行している可能性もあるので、削除ではなくステータスの確認のみを行います。

チェック用のリソースを作成するCloudFormationテンプレート

先にチェック用のリソースを作成するテンプレートを載せておきます。

AWSTemplateFormatVersion: '2010-09-09'
Description: 'AWS CloudFormation template for deploying the SageMaker Spaces status checker and SageMaker Endpoint checker Lambda function and EventBridge rule'

Parameters:
  ScheduleExpression:
    Type: String
    Default: 'cron(0 9 * * ? *)'
    Description: 'The schedule expression for running the Lambda function'
  EmailAddress:
    Type: String
    Default: 'example@example.com'
    Description: 'The email address to send notifications to'

Resources:
  SageMakerResourceCheckerLambda:
    Type: 'AWS::Lambda::Function'
    Properties:
      FunctionName: 'SageMakerResourceChecker'
      Handler: 'index.lambda_handler'
      Runtime: 'python3.12'
      Timeout: 10
      Role: !GetAtt LambdaExecutionRole.Arn
      Environment:
        Variables:
          APP_STATUS_TOPIC_ARN: !Ref SageMakerAppStatusTopic
          ENDPOINT_STATUS_TOPIC_ARN: !Ref SageMakerEndpointStatusTopic
      Code:
        ZipFile: |
          import boto3
          import json
          import os

          def lambda_handler(event, context):
              sagemaker_client = boto3.client('sagemaker')
              sns_client = boto3.client('sns')

              # Check SageMaker space status
              app_response = sagemaker_client.list_apps()
              apps = app_response['Apps']

              app_status_list = []
              for app in apps:
                  space_name = app.get('SpaceName', '')
                  app_type = app['AppType']
                  app_name = app['AppName']
                  app_status = app['Status']

                  status = '起動中' if app_status == 'InService' else '停止中'
                  app_status_list.append(f"スペース名: {space_name}\nアプリケーションタイプ: {app_type}\nアプリケーション名: {app_name}\nステータス: {status}\n---")

              app_message = "現在のSageMakerアプリケーションの状況一覧:\n\n" + "\n".join(app_status_list)

              sns_client.publish(
                  TopicArn=os.environ['APP_STATUS_TOPIC_ARN'],
                  Message=app_message,
                  Subject='SageMakerアプリケーションの状況レポート'
              )

              # Check SageMaker endpoint status
              endpoint_response = sagemaker_client.list_endpoints()
              endpoints = endpoint_response['Endpoints']

              if endpoints:
                  endpoint_status_list = []
                  for endpoint in endpoints:
                      endpoint_name = endpoint['EndpointName']
                      endpoint_status = endpoint['EndpointStatus']
                      creation_time = endpoint['CreationTime']
                      last_modified_time = endpoint['LastModifiedTime']

                      endpoint_status_list.append(f"エンドポイント名: {endpoint_name}\nステータス: {endpoint_status}\n作成時刻: {creation_time}\n最終更新時刻: {last_modified_time}\n---")

                  endpoint_message = "現在のSageMakerエンドポイントの状況一覧:\n\n" + "\n".join(endpoint_status_list)
              else:
                  endpoint_message = "現在起動中のエンドポイントはありません。"

              sns_client.publish(
                  TopicArn=os.environ['ENDPOINT_STATUS_TOPIC_ARN'],
                  Message=endpoint_message,
                  Subject='SageMakerエンドポイントの状況レポート'
              )

              return {
                  'statusCode': 200,
                  'body': 'SageMaker application and endpoint status checked and email sent successfully.'
              }

  LambdaExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: 'sts:AssumeRole'
      Policies:
        - PolicyName: 'SageMakerResourceCheckerPolicy'
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - 'sagemaker:ListApps'
                  - 'sagemaker:ListEndpoints'
                Resource: '*'
              - Effect: Allow
                Action:
                  - 'sns:Publish'
                Resource:
                  - !Ref SageMakerAppStatusTopic
                  - !Ref SageMakerEndpointStatusTopic

  SageMakerAppStatusTopic:
    Type: 'AWS::SNS::Topic'
    Properties:
      TopicName: 'SageMakerAppStatusTopic'
      Subscription:
        - Endpoint: !Ref EmailAddress
          Protocol: 'email'

  SageMakerEndpointStatusTopic:
    Type: 'AWS::SNS::Topic'
    Properties:
      TopicName: 'SageMakerEndpointStatusTopic'
      Subscription:
        - Endpoint: !Ref EmailAddress
          Protocol: 'email'

  SageMakerResourceCheckerRule:
    Type: 'AWS::Events::Rule'
    Properties:
      Name: 'SageMakerResourceCheckerRule'
      ScheduleExpression: !Ref ScheduleExpression
      State: 'ENABLED'
      Targets:
        - Arn: !GetAtt SageMakerResourceCheckerLambda.Arn
          Id: 'SageMakerResourceCheckerLambdaTarget'

  LambdaInvokePermission:
    Type: 'AWS::Lambda::Permission'
    Properties:
      FunctionName: !Ref SageMakerResourceCheckerLambda
      Action: 'lambda:InvokeFunction'
      Principal: 'events.amazonaws.com'
      SourceArn: !GetAtt SageMakerResourceCheckerRule.Arn

デプロイ方法

CloudFormationでデプロイする際に以下を聞かれるので入力してデプロイします。

停止忘れチェックの時間(cron式)
メールアドレス

その後、作成されたSNSから確認メールが2通（スペース用、エンドポイント用）届くので確認を行います。

あとは日次で以下のようなメールが届きます。

スペースのステータス通知メール

エンドポイントのステータス通知メール

なぜ作ったのか

リソース削除忘れによる不要なコストを防ぐため

SageMakerを使用していると、Studio内で起動したリソースを消し忘れてしまうことがあります。
その中でも私がよく止め忘れてしまうエンドポイントとノートブックインスタンス(本当はスペース)の停止忘れを確認するためにこの仕組みを作りました。

今回この仕組みを実装する際に、すでに誰かが作ってるだろう！と思い調べてみました。
すると、色々出てきました。
エンドポイントについては検索にヒットしたコードを参考にステータスを確認するものがすぐに実装できました。
しかし、ノートブックインスタンスのステータスを確認するのに苦戦しました。

詳細は割愛しますが、ノートブックインスタンスの自動終了や自動停止を行う内容が書かれた記事がたくさんありました。
これらを参考に、Studioで起動したノートブックインスタンスのステータスを確認する関数を作成してみました。
しかし...何を試しても起動したはずのノートブックインスタンスのステータスを確認することができませんでした。

そこでStudioでJupyterLabを触っている時に気づきました。
あれ？これノートブックインスタンスじゃなくて、スペースって書いてある、と。

最近のアップデートでStudioが新しくなり、起動時のプロセスが変わったことでノートブックインスタンスではなくスペースなるものが作成されていることを理解しました。

ちなみに去年の末にStudioが新しくなりました！

新しいStudioと古いStudioでは起動方式が違う

古いStudioでは起動時にノートブックインスタンスが起動していました。
しかし、新しいStudioは起動時にスペースというものが作成されるようです。

ということで、新旧でそれぞれStudioからアプリケーション(ノートブックとスペース)を立ち上げてみます。

SageMakerのコンソール画面を見てみましょう。
古いStudioでは起動後にアプリケーション一覧にJupyterServerが表示されています。

しかし新しいStudioで起動したアプリケーションは一覧に表示されていません。
そこでドメインの詳細から「スペースの管理」を確認してみました。

おお、ここにいるじゃないか！
ということで、新しいStudio内でアプリケーションを起動するとスペースが作成されることがわかりました。

つまり、起動プロセスが変わる＝APIも変わったということです。
当然ですが、今まで通りの方法ではノートブックインスタンスの確認はできてもスペースのステータスは分からないということです。

スペースのステータスを確認するAPIを探してみた

SpaceのAPI

そこで、Spaceの詳細を確認するAPIを探してみました。

以下のコマンドでSpaceの詳細が確認できます。

aws sagemaker describe-space --domain-id <domain-id> --space-name <space_name>

スペースを起動した状態で以下のようなレスポンスが返ってきました。

{
    "DomainId": "a-asdfghjklqwer",
    "SpaceArn": "arn:aws:sagemaker:ap-northeast-1:1234567890:space/a-asdfghjklqwer/demo-space",
    "SpaceName": "demo-space",
    "Status": "InService",
    "LastModifiedTime": "2024-03-24T03:01:36.084000+00:00",
    "CreationTime": "2024-03-24T03:01:19.382000+00:00",
    "SpaceSettings": {
        "JupyterLabAppSettings": {
            "DefaultResourceSpec": {
                "SageMakerImageArn": "arn:aws:sagemaker:ap-northeast-1:0109277902:image/sagemaker-distribution-cpu",
                "SageMakerImageVersionAlias": "1.5.2",
                "InstanceType": "ml.t3.medium"
            }
        },
        "AppType": "JupyterLab",
        "SpaceStorageSettings": {
            "EbsStorageSettings": {
                "EbsVolumeSizeInGb": 5
            }
        }
    },
    "OwnershipSettings": {
        "OwnerUserProfileName": "my-user-profile"
    },
    "SpaceSharingSettings": {
        "SharingType": "Private"
    },
    "Url": "https://fm4ipczf6wokccm.studio.ap-northeast-1.sagemaker.aws/jupyterlab/default"
}

スペースを停止して同じように確認しました。

{
    "DomainId": "a-asdfghjklqwer",
    "SpaceArn": "arn:aws:sagemaker:ap-northeast-1:1234567890:space/a-asdfghjklqwer/demo-space",
    "SpaceName": "demo-space",
    "Status": "InService",
    ...
}

あれ、Status一緒じゃん。
アプリケーションがrun状態でもstop状態でもStatusはInServiceと表示されていました。
これではスペースのステータスを確認することができません。

appのAPI

他にSpaceのステータスを確認できるものがないかと色々探しているとlist-appというAPIを発見しました。

試しにやってみました。

$ aws sagemaker list-apps

まずは起動状態で確認してみます。

{
    "Apps": [
        {
            "DomainId": "a-asdfghjklqwer",
            "SpaceName": "demo-space2",
            "AppType": "JupyterLab",
            "AppName": "default",
            "Status": "InService",
            ...
}

次に停止状態でも確認しました。

{
    "Apps": [
        {
            "DomainId": "a-asdfghjklqwer",
            "SpaceName": "demo-space2",
            "AppType": "JupyterLab",
            "AppName": "default",
            "Status": "Deleted",
            ...
}

おお！スペースの起動状態と停止状態でStatusが変わってるじゃないか！
と、いうことで、spaceではなくappのAPIでステータスの変化が確認できることがわかりました。

では、これをInServiceは起動中、Deletedは停止中と変換することでスペースの停止忘れに気づくことができます。
今回作成したCloudFormationテンプレートはこのようにしてスペースのステータスをチェックしています。

また、新しくなったStudioでも古いStudioを起動することができます。